import os
from typing import List, Dict, Any
from PyPDF2 import PdfReader
from whoosh import index
from whoosh.fields import Schema, TEXT, ID, NUMERIC
from whoosh.qparser import QueryParser
from fastmcp import FastMCP

# 初始化MCP服务
mcp = FastMCP("pdf_indexer.mcp")

PRESET_INDEX_DIRS = {"BY": "pdf_index_BY"}


def find_pdf_files(directory: str) -> List[str]:
    pdf_files = []
    for root, _, files in os.walk(directory):
        for file in files:
            if file.lower().endswith(".pdf"):
                pdf_files.append(os.path.join(root, file))
    return pdf_files


def extract_text_from_pdf(pdf_path: str) -> List[str]:
    texts = []
    try:
        reader = PdfReader(pdf_path)
        for page in reader.pages:
            text = page.extract_text() or ""
            texts.append(text)
    except Exception as e:
        print(f"提取 {pdf_path} 时出错: {e}")
    return texts


@mcp.tool()
def build_pdf_index(pdf_dir: str, index_dir: str) -> Dict[str, Any]:
    """为指定目录下所有PDF建立全文索引

    参数:
        pdf_dir: PDF文件所在目录
        index_dir: 索引存储目录

    返回:
        dict: 包含索引结果的字典
            - success: 布尔值，表示索引是否成功
            - message: 索引结果消息
            - index_dir: 索引存储目录
            - pdf_count: 处理的PDF文件数量
            - error: 错误信息(如果success为False)
    """
    result = {
        "success": False,
        "message": None,
        "index_dir": index_dir,
        "pdf_count": 0,
        "error": None
    }

    try:
        # 检查PDF目录是否存在
        if not os.path.exists(pdf_dir):
            result["error"] = f"PDF目录不存在: {pdf_dir}"
            return result

        # 创建或打开索引目录
        if not os.path.exists(index_dir):
            os.makedirs(index_dir, exist_ok=True)

        schema = Schema(
            path=ID(stored=True), page=NUMERIC(stored=True), content=TEXT(stored=True)
        )

        if not index.exists_in(index_dir):
            ix = index.create_in(index_dir, schema)
        else:
            ix = index.open_dir(index_dir)

        writer = ix.writer()
        pdf_files = find_pdf_files(pdf_dir)
        result["pdf_count"] = len(pdf_files)

        if len(pdf_files) == 0:
            result["error"] = f"未在目录中找到PDF文件: {pdf_dir}"
            return result

        for pdf in pdf_files:
            pages = extract_text_from_pdf(pdf)
            for page_num, text in enumerate(pages, 1):
                writer.add_document(path=pdf, page=page_num, content=text)

        writer.commit()
        result["success"] = True
        result["message"] = f"索引已建立，目录：{index_dir}，共{len(pdf_files)}个PDF"

    except Exception as e:
        result["error"] = f"建立索引时出错: {str(e)}"

    return result


@mcp.tool()
def search_pdf_index(
    query: str, index_key: str = "", index_dir: str = "", limit: int = 10
) -> Dict[str, Any]:
    """在指定PDF索引目录中全文检索关键词

    参数:
        query: 搜索关键词
        index_key: 预设索引目录别名
        index_dir: 索引存储目录
        limit: 返回结果数量限制

    返回:
        dict: 包含搜索结果的字典
            - success: 布尔值，表示搜索是否成功
            - results: 搜索结果列表
            - query: 搜索关键词
            - error: 错误信息(如果success为False)
    """
    result = {
        "success": False,
        "results": [],
        "query": query,
        "error": None
    }

    try:
        resolved_index_dir = None
        if index_key and index_key in PRESET_INDEX_DIRS:
            resolved_index_dir = PRESET_INDEX_DIRS[index_key]
        elif index_dir:
            resolved_index_dir = index_dir
        else:
            result["error"] = "未指定索引目录"
            return result

        if not index.exists_in(resolved_index_dir):
            result["error"] = f"索引目录不存在: {resolved_index_dir}，请先建立索引"
            return result

        ix = index.open_dir(resolved_index_dir)
        with ix.searcher() as searcher:
            parser = QueryParser("content", schema=ix.schema)
            q = parser.parse(query)
            results = searcher.search(q, limit=limit)
            for hit in results:
                result["results"].append(
                    {
                        "file": hit["path"],
                        "page": hit["page"],
                        "snippet": hit.highlights("content"),
                    }
                )

        result["success"] = True

    except Exception as e:
        result["error"] = f"搜索时出错: {str(e)}"

    return result


@mcp.tool()
def ping() -> str:
    """健康检查"""
    return "pong"


if __name__ == "__main__":
    import os

    def run_server():
        mode = os.getenv("MCP_MODE", "").lower()
        if mode == "sse":
            mcp.run(transport="sse", port=8000, path="/sse")
        elif mode == "streamable-http":
            mcp.run(transport="streamable-http", port=8000, path="/mcp")
        else:
            mcp.run()  # 默认STDIO

    run_server()
